"""
@File    : get_capcha.py
@Author  : timo.tan
@Time    : 2025/07/23 12:25 
"""

import requests
import base64
import sys
import os
from manage_dev.captcha_auto_recognizer import auto_recognize_captcha, save_captcha_model

def get_capcha():
    url = 'http://cgmgr.vrviu.com:19997/api/base/captcha'
    response = requests.post(url=url).json()
    return response['data']  # Return the complete data dictionary


def save_captcha_image(captcha_data, output_path="captcha.png"):
    """Extract and save the base64 image data to a file"""
    if 'picPath' in captcha_data:
        # Extract the base64 part (removing the "data:image/png;base64," prefix)
        img_data = captcha_data['picPath'].split(',')[1]
        # Decode the base64 data
        image_bytes = base64.b64decode(img_data)
        # Save the image to a file
        with open(output_path, 'wb') as f:
            f.write(image_bytes)
        return output_path
    return None


def login(username="timo.tan", password="Thb147258..", captcha=None, captcha_id=None):
    """使用验证码登录系统"""
    if captcha_id is None:
        captcha_data = get_capcha()
        captcha_id = captcha_data.get('captchaId')
    
    url = "http://cgmgr.vrviu.com:19997/api/base/login"
    data = {
        "username": username, 
        "password": password, 
        "captcha": captcha, 
        "captchaId": captcha_id
    }
    response = requests.post(url=url, json=data).json()
    return response


if __name__ == '__main__':
    # 获取验证码数据
    captcha_data = get_capcha()
    captcha_id = captcha_data.get('captchaId')
    print(f"验证码ID: {captcha_id}")
    
    # 保存验证码图片
    saved_path = save_captcha_image(captcha_data)
    print(f"验证码图片已保存至: {saved_path}")
    
    # 自动识别验证码（允许手动纠正）
    captcha_numbers = auto_recognize_captcha(saved_path, allow_manual_correction=True)
    
    # 打印识别结果
    if captcha_numbers:
        print(f"最终验证码结果: {captcha_numbers}")
        
        # 询问是否使用此验证码登录
        print("是否使用此验证码登录? (y/n)")
        if input().lower() == 'y':
            username = input("请输入用户名 (默认 timo.tan): ") or "timo.tan"
            password = input("请输入密码 (默认 Thb147258..): ") or "Thb147258.."
            
            # 执行登录
            login_result = login(username, password, captcha_numbers, captcha_id)
            print(f"登录结果: {login_result}")
            
            # 如果登录成功，保存验证码样本
            if login_result.get('code') == 0:
                save_captcha_model(saved_path, captcha_numbers)
    else:
        print("未能识别验证码")

